{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "o4JZ84moBKMr",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Differentiable Fluid Simulations with Φ<sub>Flow</sub>\n",
    "\n",
    "[![Google Collab Book](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/tum-pbs/PhiFlow/blob/develop/docs/Fluids_Tutorial.ipynb)\n",
    "\n",
    "This notebook steps you through setting up fluid simulations and using TensorFlow's differentiation to optimize them.\n",
    "\n",
    "Execute the cell below to install the [Φ<sub>Flow</sub> Python package from GitHub](https://github.com/tum-pbs/PhiFlow)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "Z6YoAVKebfNV",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# !pip install --quiet phiflow\n",
    "from phi.flow import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BVV1IKVqDfLl",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Setting up a Simulation\n",
    "\n",
    "Φ<sub>Flow</sub> is vectorized but object-oriented, i.e. data are represented by Python objects that internally use tensors.\n",
    "\n",
    "First, we create grids for the quantities we want to simulate. For this example, we require a velocity field and a smoke density field.\n",
    "We sample the smoke field at the cell centers and the velocity in [staggered form](https://tum-pbs.github.io/PhiFlow/Staggered_Grids.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "MA62qshHGcsk",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "smoke = CenteredGrid(0, extrapolation.BOUNDARY, x=32, y=40, bounds=Box(x=32, y=40))  # sampled at cell centers\n",
    "velocity = StaggeredGrid(0, extrapolation.ZERO, x=32, y=40, bounds=Box(x=32, y=40))  # sampled in staggered form at face centers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OTmqFIH5FBhp",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Additionally, we want to add more smoke every time step.\n",
    "We create the `INFLOW` field from a circle (2D [`Sphere`](https://tum-pbs.github.io/PhiFlow/phi/geom/index.html#phi.geom.Sphere)) which defines where hot smoke is emitted.\n",
    "Furthermore, we are interested in running the simulation for different inflow locations.\n",
    "\n",
    "Φ<sub>Flow</sub> supports data-parallell execution via *batch dimensions*.\n",
    "When a quantity has a batch dimension with size *n*, operations involving that quantity will be performed *n* times simultaneously and the result will also have that batch dimension. Here we add the batch dimension `inflow_loc`.\n",
    "\n",
    "For an overview of the dimension types, see the [documentation](https://tum-pbs.github.io/PhiFlow/Math.html#shapes)\n",
    "or watch the [introductory tutorial video](https://youtu.be/4nYwL8ZZDK8)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "WrA3IXDxv31P",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "INFLOW_LOCATION = tensor([(4, 5), (8, 5), (12, 5), (16, 5)], batch('inflow_loc'), channel(vector='x,y'))\n",
    "INFLOW = 0.6 * CenteredGrid(Sphere(center=INFLOW_LOCATION, radius=3), extrapolation.BOUNDARY, x=32, y=40, bounds=Box(x=32, y=40))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "oH7okuZUbVS-",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The created grids are instances of the class `Grid`.\n",
    "Like tensors, grids also have the `shape` attribute which lists all batch, spatial and channel dimensions.\n",
    "[Shapes in Φ<sub>Flow</sub>](https://tum-pbs.github.io/PhiFlow/Math.html#shapes) store not only the sizes of the dimensions but also their names and types."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "eK7T8Pdua5hh",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "print(f\"Smoke: {smoke.shape}\")\n",
    "print(f\"Velocity: {velocity.shape}\")\n",
    "print(f\"Inflow: {INFLOW.shape}\")\n",
    "print(f\"Inflow, spatial only: {INFLOW.shape.spatial}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1S88yUJKa86i",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The grid values can be accessed using the `values` property."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lHrsSzG2bMKq",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "print(smoke.values)\n",
    "print(velocity.values)\n",
    "print(INFLOW.values)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "R8z7PWI_d2zM",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Grids have many more properties which are documented [here](https://tum-pbs.github.io/PhiFlow/phi/field/#phi.field.Grid).\n",
    "Also note that the staggered grid has a [non-uniform shape](https://tum-pbs.github.io/PhiFlow/Math.html#non-uniform-tensors) because the number of faces is not equal to the number of cells."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dWo_m_hYKZxM",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Running the Simulation\n",
    "\n",
    "Next, let's do some physics!\n",
    "Since the initial velocity is zero, we just add the inflow and the corresponding buoyancy force.\n",
    "For the buoyancy force we use the factor `(0, 0.5)` to specify strength and direction.\n",
    "Finally, we project the velocity field to make it incompressible.\n",
    "\n",
    "Note that the `@` operator is a shorthand for resampling a field at different points. Since `smoke` is sampled at cell centers and `velocity` at face centers, this conversion is necessary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "IZdJygJAKtNl",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "smoke += INFLOW\n",
    "buoyancy_force = smoke * (0, 0.5) @ velocity\n",
    "velocity += buoyancy_force\n",
    "velocity, _ = fluid.make_incompressible(velocity, (), Solve(rank_deficiency=0))\n",
    "\n",
    "vis.plot(smoke)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gvNzcYkyLU5o",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Let's run a longer simulation!\n",
    "Now we add the transport or *advection* operations to the simulation.\n",
    "Φ<sub>Flow</sub> provides multiple algorithms for advection.\n",
    "Here we use semi-Lagrangian advection for the velocity and MacCormack advection for the smoke distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-9MEio1eK5Ic",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "trajectory = [smoke]\n",
    "for i in range(20):\n",
    "  print(i, end=' ')\n",
    "  smoke = advect.mac_cormack(smoke, velocity, dt=1) + INFLOW\n",
    "  buoyancy_force = smoke * (0, 0.5) @ velocity\n",
    "  velocity = advect.semi_lagrangian(velocity, velocity, dt=1) + buoyancy_force\n",
    "  velocity, _ = fluid.make_incompressible(velocity, (), Solve(rank_deficiency=0))\n",
    "  trajectory.append(smoke)\n",
    "trajectory = field.stack(trajectory, batch('time'))\n",
    "vis.plot(trajectory, animate='time')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rdSTbMoaS0Uz",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Obtaining Gradients\n",
    "\n",
    "The simulation we just computed was using pure NumPy so all operations were non-differentiable.\n",
    "To enable differentiability, we need to use either PyTorch, TensorFlow or Jax.\n",
    "This can be achieved by changing the import statement to [`phi.tf.flow`](https://tum-pbs.github.io/PhiFlow/phi/tf/flow.html), [`phi.torch.flow`](https://tum-pbs.github.io/PhiFlow/phi/torch/flow.html) or [`phi.jax.flow`](https://tum-pbs.github.io/PhiFlow/phi/jax/index.html), respectively.\n",
    "Tensors created after this import will be allocated using PyTorch / TensorFlow / Jax and operations on these will be executed with the corresponding backend.\n",
    "These operations can make use of a GPU through CUDA if your configuration supports it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "4vxomnHQT1y4",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# from phi.jax.flow import *\n",
    "from phi.torch.flow import *\n",
    "# from phi.tf.flow import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FgE9QbjvT7d7",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We set up the simulation as before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "mphMP0sYIOz-",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "INFLOW_LOCATION = tensor([(4, 5), (8, 5), (12, 5), (16, 5)], batch('inflow_loc'), channel(vector='x,y'))\n",
    "INFLOW = 0.6 * CenteredGrid(Sphere(center=INFLOW_LOCATION, radius=3), extrapolation.BOUNDARY, x=32, y=40, bounds=Box(x=32, y=40))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GeeUQIQjf9_b",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We can verify that tensors are now backed by TensorFlow / PyTorch / Jax."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VYZ1um3Of8_r",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "type(INFLOW.values.native(INFLOW.shape))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3mpyowRYUSS4",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Note that tensors created with NumPy will keep using NumPy/SciPy operations unless a TensorFlow tensor is also passed to the same operation.\n",
    "\n",
    "Let's look at how to get gradients from our simulation.\n",
    "Say we want to optimize the initial velocities so that all simulations arrive at a final state that is similar to the right simulation where the inflow is located at `(16, 5)`.\n",
    "\n",
    "To achieve this, we define the loss function as $L = | D(s - s_r) |^2$ where $s$ denotes the smoke density and the function $D$ diffuses the difference to smoothen the gradients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "id": "EPPzC6fTkN9n",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def simulate(smoke: CenteredGrid, velocity: StaggeredGrid):\n",
    "  for _ in range(20):\n",
    "    smoke = advect.mac_cormack(smoke, velocity, dt=1) + INFLOW\n",
    "    buoyancy_force = smoke * (0, 0.5) @ velocity\n",
    "    velocity = advect.semi_lagrangian(velocity, velocity, dt=1) + buoyancy_force\n",
    "    velocity, _ = fluid.make_incompressible(velocity, (), Solve(rank_deficiency=0))\n",
    "  loss = math.sum(field.l2_loss(diffuse.explicit(smoke - field.stop_gradient(smoke.inflow_loc[-1]), 1, 1, 10)))\n",
    "  return loss, smoke, velocity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KzHC3EGhk2Yr",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now it is important that the initial velocity has the `inflow_loc` dimension before we record the gradients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "id": "ZKJ8NElOijxE",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "initial_smoke = CenteredGrid(0, extrapolation.BOUNDARY, x=32, y=40, bounds=Box(x=32, y=40))\n",
    "initial_velocity = StaggeredGrid(math.zeros(batch(inflow_loc=4)), extrapolation.ZERO, x=32, y=40, bounds=Box(x=32, y=40))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "O7iIylnIlC3F",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Finally, we use `gradient_function()` to obtain the gradient with respect to the initial velocity. Since the velocity is the second argument to the `simulate()` function, we pass `wrt=[1]`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "id": "hucSiIvjiv55",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "sim_grad = field.gradient(simulate, wrt='velocity', get_output=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "V7ceJAc_lVIp",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The argument `get_output=False` specifies that we are not interested in the actual output of the function. By setting it to `True`, we would also get the loss value and the final simulation state.\n",
    "\n",
    "To evaluate the gradient, we simply call the gradient function with the same arguments as we would call the simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "F8GMmTsdi7Mv",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "velocity_grad = sim_grad(initial_smoke, initial_velocity)\n",
    "\n",
    "vis.plot(velocity_grad)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IGt5l4JW2Hin",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "With the gradient, we can easily perform basic gradient descent optimization.\n",
    "For more advanced optimization techniques and neural network training, see the [optimization documentation](https://tum-pbs.github.io/PhiFlow/Optimization.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "iaAgH6Y3D4bT",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "print(f\"Initial loss: {simulate(initial_smoke, initial_velocity)[0]}\")\n",
    "initial_velocity -= 0.01 * velocity_grad\n",
    "print(f\"Next loss: {simulate(initial_smoke, initial_velocity)[0]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "YJe3qiLwD6OE",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "sim_grad = field.gradient(simulate, wrt='velocity', get_output=True)\n",
    "\n",
    "for opt_step in range(4):\n",
    "  (loss, final_smoke, _v), velocity_grad = sim_grad(initial_smoke, initial_velocity)\n",
    "  print(f\"Step {opt_step}, loss: {loss}\")\n",
    "  initial_velocity -= 0.01 * velocity_grad"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ooqVxCPM8PXl",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "This notebook provided an introduction to running fluid simulations in NumPy and TensorFlow.\n",
    "It demonstrated how to obtain simulation gradients which can be used to optimize physical variables or train neural networks.\n",
    "\n",
    "The full Φ<sub>Flow</sub> documentation is available at [https://tum-pbs.github.io/PhiFlow/](https://tum-pbs.github.io/PhiFlow/).\n",
    "\n",
    "Visit the [playground](https://colab.research.google.com/drive/1zBlQbmNguRt-Vt332YvdTqlV4DBcus2S#offline=true&sandboxMode=true) to run Φ<sub>Flow</sub> code in an empty notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "Differentiable Fluid Simulations with Φ-Flow 2.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "phiflow-projects",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
